package com.gitee.ixtf.rsocket

import io.cloudevents.core.builder.CloudEventBuilder
import io.rsocket.Closeable
import io.rsocket.Payload
import io.rsocket.core.RSocketClient
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.reactive.asFlow
import kotlinx.coroutines.reactor.asFlux
import kotlinx.coroutines.reactor.awaitSingle
import kotlinx.coroutines.reactor.awaitSingleOrNull
import kotlinx.coroutines.reactor.mono
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono
import reactor.util.retry.Retry

private fun <T> Mono<T>.payload(retry: Retry? = null): Mono<T> =
    if (retry == null) this else retryWhen(retry)

private fun <T> Flux<T>.payload(retry: Retry? = null): Flux<T> =
    if (retry == null) this else retryWhen(retry)

interface IRSocketAgent : Closeable {
  suspend fun rSocketClient(): RSocketClient

  suspend fun cloudEventBuilder(): CloudEventBuilder

  suspend fun metadataPush(retry: Retry? = null, block: suspend CloudEventBuilder.() -> Payload) {
    rSocketClient()
        .metadataPush(mono { block(cloudEventBuilder()) })
        .payload(retry)
        .awaitSingleOrNull()
  }

  suspend fun fireAndForget(retry: Retry? = null, block: suspend CloudEventBuilder.() -> Payload) {
    rSocketClient()
        .fireAndForget(mono { block(cloudEventBuilder()) })
        .payload(retry)
        .awaitSingleOrNull()
  }

  suspend fun requestResponse(
      retry: Retry? = null,
      block: suspend CloudEventBuilder.() -> Payload
  ): Payload =
      rSocketClient()
          .requestResponse(mono { block(cloudEventBuilder()) })
          .payload(retry)
          .awaitSingle()

  suspend fun requestStream(
      retry: Retry? = null,
      block: suspend CloudEventBuilder.() -> Payload
  ): Flow<Payload> =
      rSocketClient().requestStream(mono { block(cloudEventBuilder()) }).payload(retry).asFlow()

  suspend fun requestChannel(
      retry: Retry? = null,
      block: suspend () -> Flow<Payload>
  ): Flow<Payload> = rSocketClient().requestChannel(block().asFlux()).payload(retry).asFlow()
}
